package com.hexinfo.dmpro.component.common.service.impl;

import cn.hutool.core.util.StrUtil;
import com.hexinfo.dmpro.component.common.entity.TableFieldsEntity;
import com.hexinfo.dmpro.component.common.enums.HiveErrorFieldNamesEnum;
import com.hexinfo.dmpro.component.common.exception.BizException;
import com.hexinfo.dmpro.component.common.service.IDatabaseQueryService;
import com.hexinfo.dmpro.component.common.util.DataSourceUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.sql.*;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
import java.util.function.Predicate;
import java.util.stream.Collectors;

@Slf4j
@Component
public class HiveQueryService implements IDatabaseQueryService {

    // 表明细数据查询批次梳理
    private final static int FETCH_SIZE = 15000;

    // 表明细数据查询超时时长（秒）
    private final static int QUERY_TIMEOUT = 7200;

    /*
     * @Method: isTableExist <br>
     * @Param: [connection, dbName, tableName] <br>
     * @Return: void <br>
     * @Description：判断指定表是否存在<br>
     * @Author： wz.li<br>
     * @Date： 2023/10/19 13:30 <br>
     * @Version： V2.0.2<br>
     */
    @Override
    public void isTableExist(Connection connection, String dbName, String tableName) {
        Statement statement = null;
        ResultSet resultSet = null;
        try {
            statement = connection.createStatement();
            String template = "SHOW TABLES IN {} LIKE '{}'";
            resultSet = statement.executeQuery(StrUtil.format(template, dbName, tableName));
            if (!resultSet.next()) {
                throw new BizException(StrUtil.format("表[{}.{}]不存在！", dbName, tableName));
            }
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("指定表判断是否存在异常: {}", e.getMessage()), e);
        } finally {
            DataSourceUtil.close(statement, resultSet);
        }
    }

    /*
     * @Method: queryTotalCount <br>
     * @Param: [connection, dbTableName, queryCriteria, queryFlag] <br>
     * @Return: long <br>
     * @Description：指定表数据总量查询<br>
     * @Author： wz.li<br>
     * @Date： 2023/10/19 13:36 <br>
     * @Version： V2.0.2<br>
     */
    @Override
    public long queryTotalCount(Connection connection, String dbTableName, String queryCriteria, String queryFlag) {
        Statement statement = null;
        ResultSet resultSet = null;
        try {
            statement = connection.createStatement();
            String template = "select count(1) from {} {}";
            resultSet = statement.executeQuery(StrUtil.format(template, dbTableName, queryCriteria));
            resultSet.next();
            return resultSet.getLong(1);
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("{} 表查询总条数异常: {}", queryFlag, e.getMessage()), e);
        } finally {
            DataSourceUtil.close(statement, resultSet);
        }
    }

    /*
     * @Method: queryTableMetadata <br>
     * @Param: [conn, dbTableName, queryFlag] <br>
     * @Return: java.util.List<com.hexinfo.dmpro.component.common.entity.TableFieldsEntity> <br>
     * @Description：指定表元数据查询<br>
     * @Author： wz.li<br>
     * @Date： 2023/10/19 13:37 <br>
     * @Version： V2.0.2<br>
     */
    @Override
    public List<TableFieldsEntity> queryTableMetadata(Connection conn, String dbTableName, String queryFlag) {
        Statement statement = null;
        ResultSet resultSet = null;
        List<TableFieldsEntity> tableFieldsEntities = new ArrayList<>();
        try {
            statement = conn.createStatement();
            String template = "DESCRIBE {} ";
            resultSet = statement.executeQuery(StrUtil.format(template, dbTableName));
            while (resultSet.next()) {
                ResultSetMetaData metaData = resultSet.getMetaData();
                if (StrUtil.isEmpty(metaData.getColumnName(1))) {
                    break;
                }
                Map<String, String> tableFieldsMap = new HashMap<>();
                for (int i = 1; i <= metaData.getColumnCount(); i++) {
                    tableFieldsMap.put(metaData.getColumnName(i), resultSet.getString(i));
                }
                tableFieldsEntities.add(toTableFields(tableFieldsMap));
            }
            return tableFieldsEntities.stream().filter(t -> !HiveErrorFieldNamesEnum.getValues().contains(t.getName()))
                    .filter(distinctByKey(TableFieldsEntity::getName))
                    .collect(Collectors.toList());
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("{} 查询源数据异常: {}", queryFlag, e.getMessage()), e);
        } finally {
            DataSourceUtil.close(statement, resultSet);
        }
    }

    /*
     * @Method: distinctByKey <br>
     * @Param: [keyExtractor] <br>
     * @Return: java.util.function.Predicate<T> <br>
     * @Description：过滤<br>
     * @Author： wz.li<br>
     * @Date： 2023/10/19 13:37 <br>
     * @Version： V2.0.2<br>
     */
    private static <T> Predicate<T> distinctByKey(Function<? super T, Object> keyExtractor) {
        Map<Object, Boolean> seen = new ConcurrentHashMap<>();
        return object -> seen.putIfAbsent(keyExtractor.apply(object), Boolean.TRUE) == null;
    }

    /*
     * @Method: queryTableDetailedData <br>
     * @Param: [conn, dbTableName, queryCriteria, queryFlag] <br>
     * @Return: java.sql.ResultSet <br>
     * @Description：指定表明细数据查询<br>
     * @Author： wz.li<br>
     * @Date： 2023/10/19 13:37 <br>
     * @Version： V2.0.2<br>
     */
    @Override
    public ResultSet queryTableDetailedData(Connection conn, String dbTableName, String queryCriteria, String queryFlag) {
        Statement statement;
        try {
            statement = conn.createStatement();
            statement.setFetchSize(FETCH_SIZE);
            statement.setQueryTimeout(QUERY_TIMEOUT);
            String template = "select * from {} {}";
            String baseSql = StrUtil.format(template, dbTableName, queryCriteria);
            return statement.executeQuery(baseSql);
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("{} 表明细数据查询异常: {}", queryFlag, e.getMessage()), e);
        }
    }

    /*
     * @Method: getTableList <br>
     * @Param: [conn, dbName] <br>
     * @Return: java.util.List<java.lang.String> <br>
     * @Description：获取指定库的表集合<br>
     * @Author： wz.li<br>
     * @Date： 2023/10/19 13:37 <br>
     * @Version： V2.0.2<br>
     */
    public List<String> getTableList(Connection conn, String dbName) {
        Statement statement;
        List<String> result = new ArrayList<>();
        try {
            statement = conn.createStatement();
            statement.execute(StrUtil.format("use {}", dbName));
            ResultSet rs = statement.executeQuery(StrUtil.format("show tables"));
            while (rs.next()) {
                result.add(rs.getString(1));
            }
            statement.close();
            conn.close();
            return result;
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("查询库下所有表异常", e.getMessage()), e);
        }
    }

    /*
     * @Method: queryTablePartitionList <br>
     * @Param: [conn, dbName, tableName] <br>
     * @Return: java.util.List<java.lang.String> <br>
     * @Description：查询指定库表的分区字段集合<br>
     * @Author： wz.li<br>
     * @Date： 2023/10/19 13:38 <br>
     * @Version： V2.0.2<br>
     */
    @Override
    public List<String> queryTablePartitionList(Connection conn, String dbName, String tableName) {
        Statement statement = null;
        ResultSet resultSet = null;
        Set<String> tablePartition = new HashSet<>();
        try {
            statement = conn.createStatement();
            String template = "DESCRIBE  {}.{} ";
            resultSet = statement.executeQuery(StrUtil.format(template, dbName, tableName));
            int flag = 0;
            while (resultSet.next()) {
                String partition = resultSet.getString(1);
                if (StrUtil.isEmpty(partition)) {
                    continue;
                }
                if (flag == 0 && partition.equals("# Partition Information")) {
                    flag = 1;
                } else if (flag == 1 && partition.charAt(0) != '#') {
                    tablePartition.add(partition);
                }
            }
            return tablePartition.stream().collect(Collectors.toList());
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("获取表{}.{}分区字段异常: {}", dbName, tableName, e.getMessage()), e);
        } finally {
            DataSourceUtil.close(statement, resultSet);
        }
    }

    /*
     * @Method: truncateTable <br>
     * @Param: [conn, dbName, tableName] <br>
     * @Return: void <br>
     * @Description：删除指定表数据<br>
     * @Author： wz.li<br>
     * @Date： 2023/10/12 18:10 <br>
     * @Version： V2.0.2<br>
     */
    @Override
    public void truncateTable(Connection conn, String dbName, String tableName) {
        Statement statement = null;
        try {
            statement = conn.createStatement();
            String template = "TRUNCATE TABLE IF EXISTS {}.{}";
            statement.execute(StrUtil.format(template, dbName, tableName));
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("{}.{}表TRUNCATE异常: {}", dbName, tableName, e.getMessage()), e);
        } finally {
            DataSourceUtil.close(statement);
        }
    }

    /*
     * @Method: dropTable <br>
     * @Param: [conn, dbName, tableName] <br>
     * @Return: void <br>
     * @Description：删除指定表的数据和表结构<br>
     * @Author： wz.li<br>
     * @Date： 2023/10/12 18:09 <br>
     * @Version： V2.0.2<br>
     */
    @Override
    public void dropTable(Connection conn, String dbName, String tableName) {
        Statement statement = null;
        try {
            statement = conn.createStatement();
            String template = "DROP TABLE IF EXISTS {}.{}";
            statement.execute(StrUtil.format(template, dbName, tableName));
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("{}.{}表DROP异常: {}", dbName, tableName, e.getMessage()), e);
        } finally {
            DataSourceUtil.close(statement);
        }
    }

    /*
     * @Method: getCreateTableSql <br>
     * @Param: [conn, dbName, tableName] <br>
     * @Return: java.lang.String <br>
     * @Description：获取建表语句<br>
     * @Author： wz.li<br>
     * @Date： 2023/11/7 19:00 <br>
     * @Version： V2.0.2<br>
     */
    @Override
    public String getCreateTableSql(Connection conn, String dbName, String tableName) {
        Statement statement = null;
        ResultSet resultSet = null;
        try {
            statement = conn.createStatement();
            String template = "SHOW CREATE TABLE {}.{}";
            resultSet = statement.executeQuery(StrUtil.format(template, dbName, tableName));
            String createTableSql = "";
            while (resultSet.next()) {
                String createTableSqlParagraph = resultSet.getString(1);
                createTableSql += createTableSqlParagraph;
            }
            if (createTableSql.contains("LOCATION")) {
                return createTableSql.substring(0, createTableSql.indexOf("LOCATION")) + ";";
            } else {
                return createTableSql;
            }
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("获取{}.{}表建表SQL异常: {}", dbName, tableName, e.getMessage()), e);
        } finally {
            DataSourceUtil.close(statement, resultSet);
        }
    }

    /*
     * @Method: delTablePartition <br>
     * @Param: [conn, dbName, tableName, partitionName, partitionVal] <br>
     * @Return: void <br>
     * @Description：删除指定表的指定分区<br>
     * @Author： wz.li<br>
     * @Date： 2023/10/12 18:09 <br>
     * @Version： V2.0.2<br>
     */
    @Override
    public void delTablePartition(Connection conn, String dbName, String tableName, String partitionName, String partitionVal) {
        Statement statement = null;
        try {
            statement = conn.createStatement();
            String template = "ALTER TABLE {}.{} DROP PARTITION ({}='{}');";
            statement.execute(StrUtil.format(template, dbName, tableName, partitionName, partitionVal));
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("{}.{}表删除分区{}={}异常: {}", dbName, tableName, partitionName, partitionVal, e.getMessage()), e);
        } finally {
            DataSourceUtil.close(statement);
        }
    }

    /*
     * @Method: executeQuery <br>
     * @Param: [conn, querySql] <br>
     * @Return: java.sql.ResultSet <br>
     * @Description：执行指定查询SQL<br>
     * @Author： wz.li<br>
     * @Date： 2023/11/20 14:47 <br>
     * @Version： V2.0.2<br>
     */
    @Override
    public ResultSet executeQuery(Connection conn, String querySql) {
        Statement statement = null;
        try {
            statement = conn.createStatement();
            return statement.executeQuery(querySql);
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("执行指定SQL【{}】异常: {}", querySql, e.getMessage()), e);
        } finally {
            DataSourceUtil.close(statement);
        }
    }

    public List<String> getDbList(Connection conn) {
        Statement statement;
        List<String> result = new ArrayList<>();
        try {
            statement = conn.createStatement();
            ResultSet rs = statement.executeQuery("SHOW DATABASES");
            while (rs.next()) {
                result.add(rs.getString(1));
            }
            statement.close();
            conn.close();
            return result;
        } catch (Exception e) {
            throw new RuntimeException(StrUtil.format("查询所有库异常", e.getMessage()), e);
        }
    }

    private TableFieldsEntity toTableFields(Map<String, String> tableFieldsMap) {
        if (tableFieldsMap == null) {
            return null;
        }
        TableFieldsEntity tableFieldsEntity = new TableFieldsEntity();
        if (tableFieldsMap.containsKey("col_name")) {
            tableFieldsEntity.setName(tableFieldsMap.get("col_name"));
        }
        if (tableFieldsMap.containsKey("data_type")) {
            tableFieldsEntity.setType(tableFieldsMap.get("data_type"));
        }
        if (tableFieldsMap.containsKey("comment")) {
            tableFieldsEntity.setComment(tableFieldsMap.get("comment"));
        }
        return tableFieldsEntity;
    }

    @Override
    public Map<String, String> queryTableStatisticsData(Connection conn, String dbTableName, String queryCriteria, String queryFlag) {
        return null;
    }

    @Override
    public List queryTableData(Connection conn, String querySql) {
        return null;
    }

}
